#include <assert.h>
#include <vector>
#include <iostream>
#include <cuda_fp16.h>
#include "yololayer_dynamic.h"
#include "cuda_utils.h"
// #include "device_functions.h"

namespace Tn
{
    template<typename T> 
    void write(char*& buffer, const T& val)
    {
        *reinterpret_cast<T*>(buffer) = val;
        buffer += sizeof(T);
    }

    template<typename T> 
    void read(const char*& buffer, T& val)
    {
        val = *reinterpret_cast<const T*>(buffer);
        buffer += sizeof(T);
    }
}

using namespace Yolo;

namespace nvinfer1
{
    YoloLayerDynamicPlugin::YoloLayerDynamicPlugin(int classCount, int maxOut, const std::vector<Yolo::YoloKernel>& vYoloKernel)
    {
        // TODO: 确定合理的构造函数输入！

        mClassCount = classCount;
        // mYoloV5NetWidth = netWidth;
        // mYoloV5NetHeight = netHeight;
        mYoloV5NetWidth = 0;
        mYoloV5NetHeight = 0;
        mMaxOutObject = maxOut;
        mYoloKernel = vYoloKernel;
        mKernelCount = vYoloKernel.size();

        CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*)));
        size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2;
        for (int ii = 0; ii < mKernelCount; ii++)
        {
            CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen));
            const auto& yolo = mYoloKernel[ii];
            CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice));
        }
    }

    YoloLayerDynamicPlugin::~YoloLayerDynamicPlugin()
    {
        for (int ii = 0; ii < mKernelCount; ii++)
        {
            CUDA_CHECK(cudaFree(mAnchor[ii]));
        }
        CUDA_CHECK(cudaFreeHost(mAnchor));
    }

    YoloLayerDynamicPlugin::YoloLayerDynamicPlugin(const void* data, size_t length)
    {
        using namespace Tn;
        const char *d = reinterpret_cast<const char *>(data), *a = d;
        read(d, mClassCount);
        read(d, mThreadCount);
        read(d, mKernelCount);
        read(d, mYoloV5NetWidth);
        read(d, mYoloV5NetHeight);
        read(d, mMaxOutObject);
        mYoloKernel.resize(mKernelCount);
        auto kernelSize = mKernelCount * sizeof(YoloKernel);
        memcpy(mYoloKernel.data(), d, kernelSize);
        d += kernelSize;
        CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*)));
        size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2;
        for (int ii = 0; ii < mKernelCount; ii++)
        {
            CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen));
            const auto& yolo = mYoloKernel[ii];
            CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice));
        }
        assert(d == a + length);
    }

    int YoloLayerDynamicPlugin::getNbOutputs() const
    {
        return 1;
    }

    void YoloLayerDynamicPlugin::serialize(void* buffer) const
    {
        using namespace Tn;
        char* d = static_cast<char*>(buffer), *a = d;
        write(d, mClassCount);
        write(d, mThreadCount);
        write(d, mKernelCount);
        write(d, mYoloV5NetWidth);
        write(d, mYoloV5NetHeight);
        write(d, mMaxOutObject);
        auto kernelSize = mKernelCount * sizeof(YoloKernel);
        memcpy(d, mYoloKernel.data(), kernelSize);
        d += kernelSize;

        assert(d == a + getSerializationSize());
    }

    size_t YoloLayerDynamicPlugin::getSerializationSize() const
    {
        return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size() + sizeof(mYoloV5NetWidth) + sizeof(mYoloV5NetHeight) + sizeof(mMaxOutObject);
    }

    int YoloLayerDynamicPlugin::initialize()
    {
        return 0;
    }

    DimsExprs YoloLayerDynamicPlugin::getOutputDimensions(
            int outputIndex, const DimsExprs* inputs, int nbInputs, IExprBuilder& exprBuilder)
    {
        // the first input is the resized tensor, the second input provide the output shape information        
        //output the result to channel
        // int totalsize = mMaxOutObject * sizeof(Detection) / sizeof(float);
        // return Dims3(totalsize + 1, 1, 1);
        int single_batch_totalsize = mMaxOutObject * sizeof(Detection) / sizeof(float) + 1;

        DimsExprs batchsize;
        batchsize.nbDims = 1;
        batchsize.d[0] = inputs[0].d[0];

        DimsExprs totalsize;
        totalsize.nbDims = 1;
        totalsize.d[0] = exprBuilder.constant(single_batch_totalsize);

        DimsExprs out_dim;
        out_dim.nbDims = 3;
        out_dim.d[0] = exprBuilder.operation(DimensionOperation::kPROD, *batchsize.d[0], *totalsize.d[0]);
        out_dim.d[1] = exprBuilder.constant(1);
        out_dim.d[2] = exprBuilder.constant(1);

        return out_dim;
    }

    // Set plugin namespace
    void YoloLayerDynamicPlugin::setPluginNamespace(const char* pluginNamespace)
    {
        mPluginNamespace = pluginNamespace;
    }

    const char* YoloLayerDynamicPlugin::getPluginNamespace() const
    {
        return mPluginNamespace;
    }

    // Return the DataType of the plugin output at the requested index
    DataType YoloLayerDynamicPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const
    {
        return inputTypes[0];
        //return DataType::kFLOAT;
    }

    void YoloLayerDynamicPlugin::configurePlugin(const DynamicPluginTensorDesc* in, int nbInputs, const DynamicPluginTensorDesc* out, int nbOutputs)
    {       
    }

    bool YoloLayerDynamicPlugin::supportsFormatCombination(
        int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs)
    {
        return true;
    }

    const char* YoloLayerDynamicPlugin::getPluginType() const
    {
        return "YoloLayerDynamic_TRT";
    }

    const char* YoloLayerDynamicPlugin::getPluginVersion() const
    {
        return "1";
    }

    void YoloLayerDynamicPlugin::destroy()
    {
        delete this;
    }

    // Clone the plugin
    IPluginV2DynamicExt* YoloLayerDynamicPlugin::clone() const
    {
        YoloLayerDynamicPlugin* p = new YoloLayerDynamicPlugin(mClassCount, mMaxOutObject, mYoloKernel);
        p->setPluginNamespace(mPluginNamespace);
        return p;
    }

    //
    size_t YoloLayerDynamicPlugin::getWorkspaceSize(
        const PluginTensorDesc* inputs, int nbInputs, const PluginTensorDesc* outputs, int nbOutputs) const
    {
        return 0;
    }

    //
    __device__ float Logist_Dynamic_float(float data) { return 1.0f / (1.0f + expf(-data)); }

    __device__ float Logist_Dynamic_half(half data) { half res=1.0f; return __half2float( __hdiv(res, __hadd(res, hexp(-data))));}

    template<typename scalar_t>
    __global__ void CalDetection_Dynamic(
            nvinfer1::DataType data_type,
            const scalar_t *input, float *output, int noElements,
            const int netwidth, const int netheight, int maxoutobject,
            int yoloWidth, int yoloHeight, const float anchors[CHECK_COUNT * 2], int classes, int outputElem)
    {
        int idx = threadIdx.x + blockDim.x * blockIdx.x;
        if (idx >= noElements) return;

        int total_grid = yoloWidth * yoloHeight;
        int bnIdx = idx / total_grid;
        idx = idx - total_grid * bnIdx;
        int info_len_i = 5 + classes;
        const scalar_t* curInput = input + bnIdx * (info_len_i * total_grid * CHECK_COUNT);
        float box_prob=0.0f;
        for (int k = 0; k < CHECK_COUNT; ++k) {
            switch(data_type){
                case nvinfer1::DataType::kFLOAT:
                    box_prob = Logist_Dynamic_float(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]);
                    break;
                case nvinfer1::DataType::kHALF:
                    box_prob = Logist_Dynamic_half(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]);
                    break;
                default:
                    return;
            }

            if (box_prob < IGNORE_THRESH) continue;
            int class_id = 0;
            float max_cls_prob = 0.0;
            for (int i = 5; i < info_len_i; ++i) {
                float p;
                switch(data_type){
                    case nvinfer1::DataType::kFLOAT:
                        p = Logist_Dynamic_float(curInput[idx + k * info_len_i * total_grid + i * total_grid]);
                        break;
                    case nvinfer1::DataType::kHALF:
                        p = Logist_Dynamic_half(curInput[idx + k * info_len_i * total_grid + i * total_grid]);
                        break;
                    default:
                        return;
                }
                if (p > max_cls_prob) {
                    max_cls_prob = p;
                    class_id = i - 5;
                }
            }
            float *res_count = output + bnIdx * outputElem;
            int count = (int)atomicAdd(res_count, 1);
            if (count >= maxoutobject) return;
            char *data = (char*)res_count + sizeof(float) + count * sizeof(Detection);
            Detection *det = (Detection*)(data);

            int row = idx / yoloWidth;
            int col = idx % yoloWidth;

            switch(data_type){
                case nvinfer1::DataType::kFLOAT:
                    det->bbox[0] = (col - 0.5f + 2.0f * Logist_Dynamic_float(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * netwidth / yoloWidth;
                    det->bbox[1] = (row - 0.5f + 2.0f * Logist_Dynamic_float(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * netheight / yoloHeight;

                    det->bbox[2] = 2.0f * Logist_Dynamic_float(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]);
                    det->bbox[3] = 2.0f * Logist_Dynamic_float(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]);
                    break;
                case nvinfer1::DataType::kHALF:
                    det->bbox[0] = (col - 0.5f + 2.0f * Logist_Dynamic_half(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * netwidth / yoloWidth;
                    det->bbox[1] = (row - 0.5f + 2.0f * Logist_Dynamic_half(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * netheight / yoloHeight;

                    det->bbox[2] = 2.0f * Logist_Dynamic_half(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]);
                    det->bbox[3] = 2.0f * Logist_Dynamic_half(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]);
                    break;
                default:
                    return;
            }

            det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2 * k];
            det->bbox[3] = det->bbox[3] * det->bbox[3] * anchors[2 * k + 1];
            det->conf = box_prob * max_cls_prob;
            det->class_id = class_id;
        }
    }

    template<typename scalar_t>
    void CalDetection_Dynamic_(
            nvinfer1::DataType data_type, const scalar_t *input, float *output, int noElements,
            const int netwidth, const int netheight, int maxoutobject, int mThreadCount,
            int yoloWidth, int yoloHeight, const float anchors[CHECK_COUNT * 2], int classes,
            int outputElem, cudaStream_t stream){
        CalDetection_Dynamic<scalar_t> << < (noElements + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> >
        (data_type, input, output, noElements, netwidth, netheight, maxoutobject, yoloWidth, yoloHeight, anchors, classes, outputElem);
    }

    template void CalDetection_Dynamic_<float>(
            nvinfer1::DataType data_type, const float *input, float *output, int noElements,
            const int netwidth, const int netheight, int maxoutobject, int mThreadCount,
            int yoloWidth, int yoloHeight, const float anchors[CHECK_COUNT * 2], int classes,
            int outputElem, cudaStream_t stream);
    template void CalDetection_Dynamic_<half>(
            nvinfer1::DataType data_type, const half *input, float *output, int noElements,
            const int netwidth, const int netheight, int maxoutobject, int mThreadCount,
            int yoloWidth, int yoloHeight, const float anchors[CHECK_COUNT * 2], int classes,
            int outputElem, cudaStream_t stream);

    void YoloLayerDynamicPlugin::forwardGpu(nvinfer1::DataType data_type, const void* const* inputs, float *output, cudaStream_t stream, int batchSize)
    {
        int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float);
        for (int idx = 0; idx < batchSize; ++idx) {
            CUDA_CHECK(cudaMemset(output + idx * outputElem, 0, sizeof(float)));
        }
        int numElem = 0;
        for (unsigned int i = 0; i < mYoloKernel.size(); ++i) {
            const auto& yolo = mYoloKernel[i];
            numElem = yolo.width * yolo.height * batchSize;
            if (numElem < mThreadCount) mThreadCount = numElem;
            switch (data_type) {
                case nvinfer1::DataType::kFLOAT:
                    CalDetection_Dynamic_<float>(data_type, (const float *)inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, mThreadCount, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem, stream);
                    break;
                case nvinfer1::DataType::kHALF:
                    CalDetection_Dynamic_<half>(data_type, (const half *)inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, mThreadCount, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem, stream);
                    break;
            }
        }
    }

    int YoloLayerDynamicPlugin::enqueue(
        const PluginTensorDesc *inputDesc, 
        const PluginTensorDesc *outputDesc, 
        const void *const *inputs, 
        void *const *outputs, 
        void *workspace, 
        cudaStream_t stream)
    {
        int iYoloKernelSize = mYoloKernel.size();

        // Method 2: get input image shape information from dims struct of input image tensor
        Dims test_size = inputDesc[iYoloKernelSize].dims;
        int input_img_b = test_size.d[0]; // batchsize
        int input_img_c = test_size.d[1]; // channel
        int input_img_h = test_size.d[2]; // height
        int input_img_w = test_size.d[3]; // width

        int scale = 8;
        for (int i = 0; i < iYoloKernelSize; ++i)
        {
            auto& kernel = mYoloKernel[i];
            kernel.width = input_img_w / scale;
            kernel.height = input_img_h / scale;
            scale *= 2;
        }
        mYoloV5NetWidth = input_img_w;
        mYoloV5NetHeight = input_img_h;

        forwardGpu(inputDesc[0].type, inputs, (float *)outputs[0], stream, input_img_b);
        return 0;
    }

    //////////////////////////////////////////////////////////////////////////////////
    PluginFieldCollection YoloLayerDynamicPluginCreator::mFC{};
    std::vector<PluginField> YoloLayerDynamicPluginCreator::mPluginAttributes;

    YoloLayerDynamicPluginCreator::YoloLayerDynamicPluginCreator()
    {
        mPluginAttributes.clear();

        mFC.nbFields = mPluginAttributes.size();
        mFC.fields = mPluginAttributes.data();
    }

    YoloLayerDynamicPluginCreator::~YoloLayerDynamicPluginCreator()
    {
    }

    const char* YoloLayerDynamicPluginCreator::getPluginName() const
    {
        return "YoloLayerDynamic_TRT";
    }

    const char* YoloLayerDynamicPluginCreator::getPluginVersion() const
    {
        return "1";
    }

    const PluginFieldCollection* YoloLayerDynamicPluginCreator::getFieldNames()
    {
        return &mFC;
    }

    IPluginV2DynamicExt* YoloLayerDynamicPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
    {
        assert(fc->nbFields == 2);
        assert(strcmp(fc->fields[0].name, "netinfo") == 0);
        assert(strcmp(fc->fields[1].name, "kernels") == 0);
        int *p_netinfo = (int*)(fc->fields[0].data);
        int class_count = p_netinfo[0];
        int max_output_object_count = p_netinfo[1];
        std::vector<Yolo::YoloKernel> kernels(fc->fields[1].length);
        memcpy(&kernels[0], fc->fields[1].data, kernels.size() * sizeof(Yolo::YoloKernel));
        YoloLayerDynamicPlugin* obj = new YoloLayerDynamicPlugin(class_count, max_output_object_count, kernels);
        obj->setPluginNamespace(mNamespace.c_str());
        return obj;
    }

    IPluginV2DynamicExt* YoloLayerDynamicPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
    {
        // This object will be deleted when the network is destroyed, which will
        // call YoloLayerDynamicPlugin::destroy()
        YoloLayerDynamicPlugin* obj = new YoloLayerDynamicPlugin(serialData, serialLength);
        obj->setPluginNamespace(mNamespace.c_str());
        return obj;
    }


}

